#include "YUVItem.h"
#include "YUVTexture.h"
#include "tool.h"
#include <QByteArray>
#include <QFile>
#include <QQuickWindow>
#include <QSGGeometry>
#include <QSGGeometryNode>
#include <QSGMaterialShader>
#include <QSGTexture>
#include <QSGTextureMaterial>
#include <QTimer>
#include <QUrl>
#include <memory>
#include <private/qrhi_p.h>

class YUVShader : public QSGMaterialShader
{
public:
	YUVShader()
	{
		auto vs = getResPath(qmlPath + QStringLiteral("yuv.vert.qsb"));
		auto fs = getResPath(qmlPath + QStringLiteral("yuv.frag.qsb"));

		setShaderFileName(VertexStage, vs);
		setShaderFileName(FragmentStage, fs);
	}
	virtual bool updateUniformData(RenderState& state, QSGMaterial* newMaterial, QSGMaterial* oldMaterial) override;
	virtual void updateSampledImage(RenderState& state, int binding, QSGTexture** texture, QSGMaterial* newMaterial, QSGMaterial* oldMaterial) override;
};

class YUVMaterial : public QSGMaterial
{
public:
	YUVMaterial(int w, int h)
		: mW(w)
		, mH(h)
	{
		mTextureY = new YUVTexture;
		mTextureU = new YUVTexture;
		mTextureV = new YUVTexture;

		setFlag(Blending, false);
	}
	~YUVMaterial()
	{
		delete mTextureY;
		delete mTextureU;
		delete mTextureV;
	}

	int			mW		  = 0;
	int			mH		  = 0;
	YUVTexture* mTextureY = nullptr;
	YUVTexture* mTextureU = nullptr;
	YUVTexture* mTextureV = nullptr;
	QByteArray	mDataY, mDataU, mDataV;
	bool		mYuvDirty = false;

public:
	virtual QSGMaterialType* type() const override
	{
		static QSGMaterialType t;
		return &t;
	}
	virtual QSGMaterialShader* createShader(QSGRendererInterface::RenderMode renderMode) const override
	{
		return new YUVShader();
	}
	virtual int compare(const QSGMaterial* o) const override
	{
		Q_ASSERT(o && type() == o->type());
		const auto* other = static_cast<const YUVMaterial*>(o);
		if (other == this)
		{
			return 0;
		}
		qint64 diff = 0;
		do
		{
			diff = mTextureY->comparisonKey() - other->mTextureY->comparisonKey();
			if (0 != diff)
			{
				break;
			}
			diff = mTextureU->comparisonKey() - other->mTextureU->comparisonKey();
			if (0 != diff)
			{
				break;
			}
			diff = mTextureV->comparisonKey() - other->mTextureV->comparisonKey();
			if (0 != diff)
			{
				break;
			}
		} while (0);
		return diff < 0 ? -1 : (diff > 0 ? 1 : 0);
	}
	void updateYUV(const QByteArray& y, const QByteArray& u, const QByteArray& v)
	{
		mDataY	  = y;
		mDataU	  = u;
		mDataV	  = v;
		mYuvDirty = true;
	}
	void updateTexture(QRhi* rhi, QRhiResourceUpdateBatch* resourceUpdates)
	{
		if (!mYuvDirty)
		{
			return;
		}
		mTextureY->setData(QRhiTexture::Format::R8, QSize(mW, mH), mDataY);
		mTextureY->commitTextureOperations(rhi, resourceUpdates);

		mTextureU->setData(QRhiTexture::Format::R8, QSize(mW / 2, mH / 2), mDataU);
		mTextureU->commitTextureOperations(rhi, resourceUpdates);

		mTextureV->setData(QRhiTexture::Format::R8, QSize(mW / 2, mH / 2), mDataV);
		mTextureV->commitTextureOperations(rhi, resourceUpdates);
	}
};

bool YUVShader::updateUniformData(RenderState& state, QSGMaterial* newMaterial, QSGMaterial* oldMaterial)
{
	bool		changed = false;
	QByteArray* buf		= state.uniformData();
	Q_ASSERT(buf->size() >= 68);
	if (state.isMatrixDirty())
	{
		const auto& m = state.combinedMatrix();
		memcpy(buf->data(), m.constData(), 64);
		changed = true;
	}
	if (state.isOpacityDirty())
	{
		const float opacity = state.opacity();
		memcpy(buf->data() + 64, &opacity, 4);
		changed = true;
	}
	auto material = static_cast<YUVMaterial*>(newMaterial);
	material->updateTexture(state.rhi(), state.resourceUpdateBatch());

	return changed;
}
void YUVShader::updateSampledImage(RenderState& state, int binding, QSGTexture** texture, QSGMaterial* newMaterial, QSGMaterial* oldMaterial)
{
	(void)state;
	(void)binding;

	auto material = static_cast<YUVMaterial*>(newMaterial);

	if (binding == 1)
	{
		// material->mTextureY->commitTextureOperations(state.rhi(), state.resourceUpdateBatch());
		*texture = material->mTextureY;
	}
	else if (binding == 2)
	{
		// material->mTextureU->commitTextureOperations(state.rhi(), state.resourceUpdateBatch());
		*texture = material->mTextureU;
	}
	else if (binding == 3)
	{
		// material->mTextureV->commitTextureOperations(state.rhi(), state.resourceUpdateBatch());
		*texture = material->mTextureV;
	}
}

class YUVNode : public QSGGeometryNode
{
public:
	YUVNode(int w, int h)
	{
		mMaterial = new YUVMaterial(w, h);
		setMaterial(mMaterial);
		setFlag(OwnsMaterial);

		mGeo = new QSGGeometry(QSGGeometry::defaultAttributes_TexturedPoint2D(), 4);
		QSGGeometry::updateTexturedRectGeometry(mGeo, QRect(), QRect());
		setGeometry(mGeo);
		setFlag(OwnsGeometry);
	}
	void setRect(const QRectF& bounds)
	{
		QSGGeometry::updateTexturedRectGeometry(mGeo, bounds, QRectF(0, 0, 1, 1));
		markDirty(QSGNode::DirtyGeometry);
	}
	void setYUV(const QByteArray& y, const QByteArray& u, const QByteArray& v)
	{
		mMaterial->updateYUV(y, u, v);
		markDirty(QSGNode::DirtyMaterial);
	}

private:
	QSGGeometry* mGeo;
	YUVMaterial* mMaterial;
};

YUVItem::YUVItem(QQuickItem* parent)
	: Super(parent)
{
	setFlag(Super::ItemHasContents);

	connect(
		this,
		&QQuickItem::xChanged,
		this,
		[this]()
		{
			mGeometryChanged = true;
			update();
		});
	connect(
		this,
		&QQuickItem::yChanged,
		this,
		[this]()
		{
			mGeometryChanged = true;
			update();
		});
	connect(
		this,
		&QQuickItem::widthChanged,
		this,
		[this]()
		{
			mGeometryChanged = true;
			update();
		});
	connect(
		this,
		&QQuickItem::heightChanged,
		this,
		[this]()
		{
			mGeometryChanged = true;
			update();
		});

	mTimer = new QTimer(this);
	connect(mTimer, &QTimer::timeout, this, &YUVItem::onTimeout);

	connect(this, &YUVItem::sourceChanged, this, &YUVItem::onSourceChanged);
}

YUVItem::~YUVItem() { }

QSGNode* YUVItem::updatePaintNode(QSGNode* oldNode, UpdatePaintNodeData* data)
{
	(void)data;
	auto node = static_cast<YUVNode*>(oldNode);
	if (!node)
	{
		node = new YUVNode(mDataWidth, mDataHeight);
	}
	if (mGeometryChanged)
	{
		node->setRect(boundingRect());
		mGeometryChanged = false;
	}
	if (mYUVChanged)
	{
		node->setYUV(mDataY, mDataU, mDataV);
		mYUVChanged = false;
	}
	return node;
}

void YUVItem::onTimeout()
{
	if (mYUVChanged)
	{
		return;
	}
	if (mDataOffset >= mData.size())
	{
		return;
	}
	if (mDataWidth * mDataHeight <= 0)
	{
		return;
	}
	int frameSize = mDataWidth * mDataHeight + (mDataWidth * mDataHeight) / 2;
	if (mDataOffset + frameSize >= mData.size())
	{
		return;
	}
	auto frameData = mData.mid(mDataOffset, frameSize);
	mDataY		   = frameData.left(mDataWidth * mDataHeight);
	mDataU		   = frameData.mid(mDataWidth * mDataHeight, mDataWidth * mDataHeight / 4);
	mDataV		   = frameData.right(mDataWidth * mDataHeight / 4);

	mDataOffset += frameSize;
	mYUVChanged = true;
	update();
}

void YUVItem::onSourceChanged()
{
	auto  src = getResPath(source().toString());
	QFile file(src);
	if (!file.open(QFile::ReadOnly))
	{
		qWarning() << "open file failed" << file.errorString();
		return;
	}
	mTimer->stop();
	mData = file.readAll();
	file.close();
	mDataWidth	= 352;
	mDataHeight = 288;
	mTimer->start(1000 / 30);
}
